#!/usr/bin/env python3
# Author: Armit
# Create Time: 周一 2025/07/21

# 观察数据集统计特征 (文本句式)
# NOTE: 训练集和测试集都是模板句式生成出来的，直接爆破句式可以吗 :(

import json
from pathlib import Path
from time import time
from pprint import pprint as pp
from re import compile as Regex, Pattern
from typing import List, Tuple, Union

BASE_PATH = Path(__file__).parent
DATA_PATH = BASE_PATH / 'data'
OUT_PATH = BASE_PATH / 'out'
OUT_PATH.mkdir(exist_ok=True)

PLT_NAME = 'plant'
CLS_NAME = 'class'
ENV_NAME = 'environment'
ARE_NAME = 'area'
ALT_NAME = 'altitude'

TEXT_TMPL_LIST_RAW: List[Union[Pattern, Tuple[str]]] = [
  (PLT_NAME, CLS_NAME, ENV_NAME),
  Regex('(.*)，一种美丽的(.*)植物，点缀着(.*)。'),
  Regex('(.*)的生存状态为(.*)，其家园多选在(.*)。'),
  Regex('(.*)面临(.*)的威胁，保护其(.*)刻不容缓。'),
  Regex('识别(.*)：它是(.*)物种，通常与(.*)相伴。'),
  Regex('(.*)是一种(.*)植物，通常可以在(.*)中找到。'),
  Regex('(.*)，一个(.*)的名字，它的故事发生在(.*)。'),
  Regex('专家提示：(.*)是(.*)物种，请勿破坏其(.*)。'),
  Regex('野外记录显示，(.*)（(.*)）主要出现在(.*)。'),
  Regex('若要寻找(.*)这种(.*)植物，(.*)是关键线索。'),
  Regex('对于(.*)（(.*)），(.*)是其基因多样性的摇篮。'),
  Regex('要恢复(.*)（(.*)）的种群，必须先恢复其(.*)。'),
  Regex('文献记载，(.*)的保护级别为(.*)，常见于(.*)。'),
  Regex('(.*)，这种被归类为(.*)的物种，主要生活在(.*)。'),
  Regex('了解(.*)的生活习性：它属于(.*)，且常分布于(.*)。'),
  Regex('(.*)不仅是珍稀的(.*)植物，它对(.*)的要求也相当高。'),
  Regex('(.*)之所以被列为(.*)，部分原因在于其对(.*)的依赖。'),
  Regex('生态位研究指出，(.*)（(.*)）在(.*)中扮演特定角色。'),
  Regex('如果你对(.*)（(.*)）感兴趣，可以关注其典型的(.*)。'),
  Regex('研究表明，(.*)被列为(.*)，其理想的(.*)具有特定条件。'),

  (PLT_NAME, ARE_NAME, ALT_NAME),
  Regex('(.*)主要分布在(.*)，通常生长在(.*)的高度。'),
  Regex('(.*)的生存范围在(.*)，其垂直跨度可达(.*)。'),
  Regex('观测表明，(.*)在(.*)的种群密度随(.*)变化。'),
  Regex('(.*)的足迹延伸至(.*)，其生态位常与(.*)关联。'),
  Regex('历史记录中，(.*)曾广泛分布于(.*)的(.*)地带。'),
  Regex('(.*)在(.*)的存在证明了其对(.*)环境的适应力。'),
  Regex('(.*)，其原产地据信在(.*)，常见于(.*)的山坡。'),
  Regex('该植物(.*)的分布中心在(.*)，海拔上限为(.*)。'),
  Regex('若想研究(.*)，(.*)的(.*)区域是重要的考察点。'),
  Regex('想要一睹(.*)的风采，可以前往(.*)的(.*)山区。'),
  Regex('(.*)，这种植物见于(.*)，其生长点位于(.*)左右。'),
  Regex('专家指出，(.*)在(.*)的分布格局受(.*)影响显著。'),
  Regex('追踪(.*)的分布：它出现在(.*)，最高记录在(.*)。'),
  Regex('结合地理信息，(.*)的分布点主要落在(.*)的(.*)。'),
  Regex('(.*)的地理分布以(.*)为主，其垂直分布可达到(.*)。'),
  Regex('(.*)的自然分布区域包括(.*)，能够耐受(.*)的条件。'),
  Regex('物种(.*)不仅局限于(.*)，还能适应(.*)的高山气候。'),
  Regex('地理调查确认(.*)分布在(.*)，通常局限于(.*)范围。'),
  Regex('根据记录，(.*)遍布于(.*)，尤其在(.*)区域较为常见。'),
  Regex('该物种(.*)的分布区域主要为(.*)，最高可出现在(.*)。'),
  Regex('文献资料表明，(.*)广泛分布于(.*)，适宜生长的(.*)。'),
  Regex('若绘制(.*)的分布图，你会发现它集中于(.*)的(.*)区域。'),
  Regex('对于(.*)而言，(.*)和(.*)是定义其生境的两个重要维度。'),
  Regex('(.*)是一种区域性物种，主要见于(.*)，海拔一般在(.*)。'),
  Regex('(.*)的分布地图显示其集中在(.*)，海拔适应范围包括(.*)。'),
  Regex('研究显示，(.*)的活动范围覆盖(.*)，并且能够适应(.*)的环境。'),

  (CLS_NAME, PLT_NAME, ENV_NAME),
  Regex('(.*)的(.*)展现了对(.*)独特的适应性。'),
  Regex('调查发现，(.*)的(.*)种群多集中于(.*)。'),
  Regex('(.*)的(.*)依赖于特定的(.*)来完成生命周期。'),
  Regex('属于(.*)的(.*)，其生态习性使其倾向于(.*)。'),
  Regex('作为(.*)物种，(.*)偏爱(.*)作为其主要栖息地。'),
  Regex('自然纪录片展示了(.*)的(.*)如何在(.*)中生存。'),
  Regex('珍稀程度达到(.*)的(.*)，与(.*)形成了共生关系。'),
  Regex('被评定为(.*)的(.*)，其身影往往与(.*)紧密相连。'),
  Regex('为了保护(.*)的(.*)，维持其赖以生存的(.*)至关重要。'),

  (ARE_NAME, ALT_NAME, PLT_NAME),
  Regex('人们曾在(.*)的(.*)观测并记录到了(.*)。'),
  Regex('在(.*)的旅行中，有幸在(.*)处见到了(.*)。'),
  Regex('在(.*)的(.*)地带，是寻找(.*)的理想区域。'),
  Regex('(.*)的(.*)提供了适合(.*)生长的独特微气候。'),
  Regex('探索(.*)时，可以在(.*)附近留意(.*)的踪迹。'),
  Regex('从(.*)到(.*)的高度，构成了(.*)的主要生活空间。'),
  Regex('(.*)以及相似的(.*)地带，是(.*)的典型分布模式。'),

  (ARE_NAME, PLT_NAME, ALT_NAME),
  Regex('在(.*)，(.*)的存在与(.*)密切相关。'),
  Regex('在(.*)境内，(.*)多生长于(.*)的山地环境。'),
  Regex('不仅在(.*)，(.*)甚至能攀升至(.*)进行生长。'),
  Regex('(.*)是(.*)的重要家园之一，其生长上限约为(.*)。'),
  Regex('(.*)的生物多样性中包含了(.*)，它适应了那里的(.*)。'),

  (ALT_NAME, PLT_NAME, ARE_NAME),
  Regex('(.*)的高度是(.*)在(.*)常见的一个生长指标。'),
  Regex('(.*)是影响(.*)在(.*)分布的一个关键生态因子。'),

  (ENV_NAME, PLT_NAME, CLS_NAME),
  Regex('(.*)区域是(.*)（(.*)）的重要庇护所。'),
  Regex('在(.*)条件下，(.*)被认定为(.*)级别。'),
  Regex('适应了(.*)的(.*)，目前的保护状况是(.*)。'),
  Regex('在(.*)中，(.*)作为(.*)物种，需要特别关注。'),
  Regex('当地的(.*)是支持(.*)（(.*)）生存的关键因素。'),
  Regex('你可以在特定的(.*)中寻找到(.*)，请注意它是(.*)的物种。'),

  (ENV_NAME, CLS_NAME, PLT_NAME),
  Regex('(.*)为(.*)的(.*)提供了必要的生存条件。'),
  Regex('探索(.*)，或许能邂逅(.*)的珍稀植物(.*)。'),
  Regex('穿越(.*)，我们寻找着(.*)植物(.*)的踪迹。'),
  Regex('生态报告强调了(.*)对(.*)植物(.*)的重要性。'),
  Regex('当环境为(.*)时，有利于(.*)植物(.*)的生长繁衍。'),
  Regex('在典型的(.*)里，我们有时能观察到(.*)级别的(.*)。'),
]
TEXT_TMPL_LIST: List[Tuple[Pattern, Tuple[str]]] = []
if 'preprocess TEXT_TMPL_LIST':
  cur_comp_names = None
  for it in TEXT_TMPL_LIST_RAW:
    if isinstance(it, tuple):
      cur_comp_names = it
    else:
      TEXT_TMPL_LIST.append((it, cur_comp_names))
  del cur_comp_names
print('>> n_tmpl:', len(TEXT_TMPL_LIST))  # 80


def stats_testset():
  fp = DATA_PATH / 'test.jsonl'
  items = []
  with open(fp, 'r', encoding='utf-8') as fh:
    for line in fh.readlines():
      it = json.loads(line.strip())
      items.append(it)
  print(f'>> file {fp}')
  print(f'>> nlen: {len(items)}')

  PREDICT = []
  PHRASES = {
    PLT_NAME: set(),
    CLS_NAME: set(),
    ENV_NAME: set(),
    ARE_NAME: set(),
    ALT_NAME: set(),
  }
  for it in items:
    text: str = it['text']
    found = False
    for tmpl, comp_names in TEXT_TMPL_LIST:
      if tmpl.match(text):
        for m in tmpl.findall(text):
          comp_vals = list(m)
          label = {}
          for n, v in zip(comp_names, comp_vals):
            PHRASES[n].add(v)
            pos_start = text.find(v)
            if n not in label: label[n] = {}
            label[n][v] = [[pos_start, pos_start + len(v) - 1]]
          PREDICT.append({'text': text, 'label': label})
        found = True
        break
    if not found:
      print(f'>> need new pattern for: {text}')

  save_fp = OUT_PATH / 'testset-phrases.json'
  print(f'>> save file: {save_fp}')
  with open(save_fp, 'w', encoding='utf-8') as fh:
    json.dump({k: sorted(v) for k, v in PHRASES.items()}, fh, indent=2, ensure_ascii=False)

  save_fp = OUT_PATH / 'predict.jsonl'
  print(f'>> save file: {save_fp}')
  with open(save_fp, 'w', encoding='utf-8') as fh:
    for pred in PREDICT:
      line = json.dumps(pred, indent=None, ensure_ascii=False)
      fh.write(line)
      fh.write('\n')


if __name__ == '__main__':
  ts_start = time()
  stats_testset()
  ts_end = time()
  print(f'>> Done {ts_end - ts_start:.3f}s')
